#质心计算过程
import numpy as np
import matplotlib.pyplot as plt

plt.figure(figsize=[12, 6])
spr = 1
spc = 2
spn = 0

# 样本集
X = np.array([[1, 2], [2, 2], [6, 8], [7, 8]], dtype=np.float64)  # m x n
#定义初始化质心
C = np.array([[1, 2], [2, 2]], dtype=np.float64)  # n_cluster x n
m = len(X)
n_cluster = len(C)

spn += 1
plt.subplot(spr, spc, spn)
plt.scatter(X[:, 0], X[:, 1], color='b')
plt.scatter(C[:, 0], C[:, 1], color='r')

iters = 5
while (iters > 0):
    iters -= 1
    B = []  # n_cluster x m
    for c in C:
        dis = np.sqrt(((X - c)**2).sum(axis=1))
        B.append(dis)
    print(B)
    min_idx = np.argmin(np.array(B).T, axis=1)
    print(min_idx)  # vector (m, )

    for i in range(n_cluster):
        C[i] = np.mean(X[min_idx==i], axis=0)

spn += 1
plt.subplot(spr, spc, spn)
plt.scatter(X[:, 0], X[:, 1], color='b')
plt.scatter(C[:, 0], C[:, 1], color='r')

plt.show()
